"""
LLM4CP数据加载器文件
处理真实空间相关性数据的加载和预处理
"""

import torch
from torch.utils.data import DataLoader, TensorDataset
import numpy as np
import scipy.io as sio
import glob
import os
from tqdm import tqdm


class SimpleLLM4CPDataLoader:
    """简化的数据加载器"""
    def __init__(self, config):
        self.config = config
        self.spatial_files = glob.glob(os.path.join(config.spatial_data_dir, "spatial_correlation_*.mat"))
        print(f"🔍 找到 {len(self.spatial_files)} 个空间相关性文件")
        
    def load_real_spatial_data(self):
        """加载真实空间相关性数据"""
        all_spatial_data = []
        
        print("📊 加载真实空间相关性数据...")
        
        for file_path in self.spatial_files[:250]:  # 加载更多数据
            try:
                mat_data = sio.loadmat(file_path)
                spatial_vec = mat_data['p']  # [1, 64]
                
                if spatial_vec.shape == (1, 64):
                    spatial_complex = spatial_vec.reshape(64)
                    
                    if np.iscomplexobj(spatial_complex):
                        real_part = np.real(spatial_complex)
                        imag_part = np.imag(spatial_complex)
                    else:
                        real_part = spatial_complex
                        imag_part = np.zeros_like(spatial_complex)
                    
                    spatial_combined = np.concatenate([real_part, imag_part])  # [128]
                    all_spatial_data.append(spatial_combined)
                    
            except Exception as e:
                continue
        
        spatial_array = np.array(all_spatial_data)
        print(f"✅ 加载 {len(all_spatial_data)} 个真实空间相关性样本: {spatial_array.shape}")
        
        return spatial_array
    
    def create_data_loaders(self):
        """创建数据加载器"""
        print("🔧 创建数据加载器...")
        
        # 加载真实空间数据
        real_spatial_data = self.load_real_spatial_data()
        n_samples = len(real_spatial_data)
        
        # 创建时间序列数据
        channel_data_list = []
        spatial_data_list = []
        
        seq_total = self.config.prev_len + self.config.pred_len
        
        for i in range(n_samples - seq_total):
            # 信道数据
            channel_seq = torch.randn(seq_total, self.config.enc_in) * 0.3
            
            # 添加与空间数据的关联
            for t in range(seq_total):
                if i + t < len(real_spatial_data):
                    spatial_influence = real_spatial_data[i + t][:self.config.enc_in] * 0.15
                    channel_seq[t] += torch.FloatTensor(spatial_influence)
            
            channel_data_list.append(channel_seq)
            
            # 空间相关性数据序列
            spatial_seq = []
            for t in range(seq_total):
                if i + t < len(real_spatial_data):
                    spatial_seq.append(real_spatial_data[i + t])
                else:
                    spatial_seq.append(real_spatial_data[-1])
            
            spatial_data_list.append(np.array(spatial_seq))
        
        # 转换为张量
        channel_data = torch.stack(channel_data_list)
        spatial_data = torch.FloatTensor(np.array(spatial_data_list))
        
        print(f"📊 数据统计:")
        print(f"   样本数量: {len(channel_data)}")
        print(f"   信道数据: {channel_data.shape}")
        print(f"   空间数据: {spatial_data.shape}")
        
        # 创建数据集
        dataset = TensorDataset(channel_data, spatial_data)
        
        # 分割训练/验证
        train_size = int(0.8 * len(dataset))
        val_size = len(dataset) - train_size
        train_subset, val_subset = torch.utils.data.random_split(dataset, [train_size, val_size])
        
        train_loader = DataLoader(train_subset, batch_size=self.config.batch_size, shuffle=True)
        val_loader = DataLoader(val_subset, batch_size=self.config.batch_size, shuffle=False)
        
        print(f"✅ 训练样本: {len(train_subset)}, 验证样本: {len(val_subset)}")
        
        return train_loader, val_loader 